from flask import Flask, request, render_template, send_file, jsonify
import os
import shutil
import time
from pdf2zh import translate
from pdf2zh.doclayout import OnnxModel
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = os.path.abspath('uploads')
app.config['DOWNLOAD_FOLDER'] = os.path.abspath('downloads')

# 确保上传和下载文件夹存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['DOWNLOAD_FOLDER'], exist_ok=True)

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # 检查是否有文件上传
        if 'file' not in request.files:
            return jsonify({"success": False, "message": "No file uploaded"}), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({"success": False, "message": "No selected file"}), 400
        
        # 清理文件名，移除路径分隔符
        filename = secure_filename(file.filename)
        
        # 保存上传的文件
        upload_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(upload_path)
        
        #确定翻译使用的目标语言
        source_language = request.form.get('source_language', 'en')  # 默认源语言为英语
        target_language = request.form.get('target_language', 'zh')  # 默认目标语言为中文
        print("source:" , source_language)
        print("target:" ,  target_language)
        # 调用PDF翻译服务
        #envs = {"OLLAMA_HOST" : "http://100.118.129.86:11434" , "OLLAMA_MODEL":"gemma2"}
        #envs = {"DEEPSEEK_API_KEY" : "sk-8a224d7874064cb9bc0e69a6394afbac" , "DEEPSEEK_MODEL":"deepseek-chat"}
        envs = {"OPENAI_BASE_URL" : "https://ark.cn-beijing.volces.com/api/v3" ,
                "OPENAI_API_KEY" : "52bb73c9-3785-4ab4-ab42-6013365609fc",
                "OPENAI_MODEL" : "ep-20250213171312-gpgk7"}
        model = OnnxModel.load_available()
        params = {
            "model": model,
            'lang_in': source_language,
            'lang_out': target_language,
            'service': 'ollama',
            'thread': 4,
        }
        # 前端工作目前还是有点不正常
        #http://100.118.129.86:11434
        
        translated_files = translate(files=[upload_path], **params)
        
        # 获取翻译后的文件路径
        file_mono, file_dual = translated_files[0]
        
        # 将翻译后的文件移动到 downloads 文件夹
        mono_file = os.path.basename(file_mono)
        dual_file = os.path.basename(file_dual)
        shutil.move(os.path.abspath(file_mono), os.path.join(app.config['DOWNLOAD_FOLDER'], mono_file))
        shutil.move(os.path.abspath(file_dual), os.path.join(app.config['DOWNLOAD_FOLDER'], dual_file))
        
        # 返回下载链接
        return jsonify({
            "success": True,
            "download_url": f"/download/{mono_file}"
        })
    
    return render_template('better-index.html')

@app.route('/download/<filename>')
def download_file(filename):
    return send_file(os.path.join(app.config['DOWNLOAD_FOLDER'], filename), as_attachment=True)
@app.route('/contact')
def contact():
    return render_template('contact.html')

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=5000,debug=True)
